package me.helong.videolive;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;
import org.springframework.web.util.UriComponents;
import org.springframework.web.util.UriComponentsBuilder;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

@Component
public class VideoTransferWebSocket extends AbstractWebSocketHandler {
    private final static Logger logger = LoggerFactory.getLogger(VideoTransferWebSocket.class);
    @Value("${videolive.ffmpeg-path}")
    private String ffmpegPath;
    @Value("${videolive.rtmp-uri}")
    private String rtmpUri;
    @Value("${videolive.play-uri}")
    private String playUri;

    @Autowired
    private WebSocketSessionDataStorage sessionDataStorage;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        UriComponents uriComponents = UriComponentsBuilder.fromUri(session.getUri()).build();
        String liveChannel = uriComponents.getQueryParams().getFirst("liveChannel");
        if(liveChannel == null || liveChannel.length() == 0) {
            session.close();
            return;
        }
        Map<String, Object> sessionData = this.sessionDataStorage.getSessionData(session);
        sessionData.put("liveChannel", liveChannel);
        session.setBinaryMessageSizeLimit(5 * 1024 * 1024);
        session.setTextMessageSizeLimit(1024 * 1024);
        Thread thread = new Thread(() -> {
            Process process = null;
            try {
                String command = ffmpegPath + " -i -  -vcodec copy -acodec aac -s vga -b:v 1.5M -g 50 -f flv " + rtmpUri + liveChannel;
                //String command = ffmpegPath + " -i - -preset veryfast -f flv rtmp://docker-host/live/" + liveChannel;
                //String command = ffmpegPath + " -i - -preset veryfast -f flv rtmp://docker-host/live/" + liveChannel;
                logger.debug("频道{}推送进程：{}", liveChannel, command);
                process = Runtime.getRuntime().exec(command);
                sessionData.put("ffmpegProcess", process);
                this.new ProcessStreamReader(session, process, process.getInputStream()).start();
                this.new ProcessStreamReader(session, process, process.getErrorStream()).start();
            } catch (Exception e) {
                try {
                    session.close();
                } catch (IOException e1) {
                    logger.error("关闭频道{}异常", liveChannel, e1);
                }
                if(process != null) {
                    process.destroy();
                }
                logger.error("开始推送视频异常，频道{}即将关闭", liveChannel, e);
            }
            try {
                Map<String, Object> result = new HashMap<>();
                result.put("code", 0);
                result.put("play url", playUri + liveChannel + ".flv");
                ObjectMapper objectMapper = new ObjectMapper();
                session.sendMessage(new TextMessage(objectMapper.writeValueAsString(result).getBytes(StandardCharsets.UTF_8)));
            } catch (IOException e) {
                process.destroy();
                try {
                    session.close();
                } catch (IOException e1) {
                    e1.printStackTrace();
                }
                logger.error("频道{}开启成功，回应客户端失败，关闭频道", liveChannel, e);
            }

            logger.info("频道{}已打开", liveChannel);
        });
        thread.setDaemon(true);
        thread.start();
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) {
        Map<String, Object> sessionData = this.sessionDataStorage.getSessionData(session);
        String liveChannel = (String) sessionData.get("liveChannel");
        Process process = (Process) sessionData.get("ffmpegProcess");
        try {
            OutputStream out = process.getOutputStream();
            out.write(message.getPayload().array());
            logger.debug("频道{}推送数据：{}", liveChannel, message.toString());
        } catch (Exception e) {
            try {
                session.close();
            } catch (IOException e1) {
                logger.error("关闭频道{}异常", liveChannel, e1);
            }
            process.destroy();
            logger.error("推送异常，关闭频道{}", liveChannel, e);
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {}

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Map<String, Object> sessionData = this.sessionDataStorage.getSessionData(session);
        String liveChannel = (String) sessionData.get("liveChannel");
        Process process = (Process) sessionData.get("ffmpegProcess");
        process.destroy();
        this.sessionDataStorage.removeSessionData(session);
        logger.info("频道{}已关闭", liveChannel);
    }

    class ProcessStreamReader extends Thread {
        private final Logger logger = LoggerFactory.getLogger(ProcessStreamReader.class);

        private String liveChannel;
        private WebSocketSession session;
        private Process process;
        private InputStream in;
        private BufferedReader reader;

        public ProcessStreamReader(WebSocketSession session, Process process, InputStream in) throws UnsupportedEncodingException {
            Map<String, Object> sessionData = sessionDataStorage.getSessionData(session);
            this.liveChannel = (String) sessionData.get("liveChannel");
            this.session = session;
            this.process = process;
            this.in = in;
            this.reader = new BufferedReader(new InputStreamReader(in, "UTF-8"));
        }

        @Override
        public void run() {
            try {
                while (process.isAlive()) {
                    String line = reader.readLine();
                    if(line != null && line.length() > 0) {
                        logger.debug("收到频道{}的ffmpeg响应：{}", liveChannel, line);
                    }
                    Thread.sleep(10);
                }
            } catch (Exception e) {
                process.destroy();
                try {
                    session.close();
                } catch (IOException e1) {
                    logger.error("关闭频道{}失败", liveChannel, e1);
                }
                logger.error("读取频道{}的ffmpeg响应信息失败，频道即将关闭", liveChannel, e);
            }
        }
    }
}
